import numpy as np
import os

def gen_gard_data_simple():
    # 生成包含正负数的输入数据
    input_x = np.random.uniform(-2.3, 4.7, [8, 2048]).astype(np.float32)
    
    # 定义Hardshrink参数 λ
    lambda_val = 0.5
    
    # 计算Hardshrink结果：f(x) = x if |x| > λ else 0
    golden = np.where(np.abs(input_x) > lambda_val, input_x, 0).astype(np.float32)

    # 确保输入输出目录存在
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # 保存文件
    input_x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")
    
    print(f"生成数据完成！")
    print(f"输入数据形状: {input_x.shape}")
    print(f"输出数据形状: {golden.shape}")
    print(f"λ = {lambda_val}")
    print(f"输入范围: [{input_x.min():.3f}, {input_x.max():.3f}]")
    print(f"输出范围: [{golden.min():.3f}, {golden.max():.3f}]")
    
    # 验证一些样本
    print("\n样本验证:")
    for i in range(3):
        x_sample = input_x[0, i]
        y_sample = golden[0, i]
        print(f"x={x_sample:.3f}, |x|={abs(x_sample):.3f}, f(x)={y_sample:.3f}")

if __name__ == "__main__":
    gen_gard_data_simple()